import torch
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
import yaml


def get_transforms():
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(brightness=0.2, contrast=0.2),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    return train_transform, val_transform


def prepare_data(config):
    train_transform, val_transform = get_transforms()

    train_dataset = datasets.ImageFolder(
        root=f"{config['data']['root_dir']}/{config['data']['train_dir']}",
        transform=train_transform
    )

    val_dataset = datasets.ImageFolder(
        root=f"{config['data']['root_dir']}/{config['data']['val_dir']}",
        transform=val_transform
    )

    train_loader = DataLoader(
        dataset=train_dataset,
        batch_size=config['data']['batch_size'],
        shuffle=True,
        num_workers=config['data']['num_workers']
    )

    val_loader = DataLoader(
        dataset=val_dataset,
        batch_size=config['data']['batch_size'],
        shuffle=False,
        num_workers=config['data']['num_workers']
    )

    return train_loader, val_loader
